# Copyright (c) Alibaba, Inc. and its affiliates.
from dataclasses import dataclass, field
from typing import List, Literal, Optional

from transformers.utils import strtobool

from swift.llm import get_model_arch
from swift.utils import get_logger

logger = get_logger()


@dataclass
class TunerArguments:
    """
    TunerArguments 是一个数据类，用于保存各种调谐器的配置。

    参数：
    target_modules（List[str]）：待调谐的目标模块列表。默认值为 ['all-linear']。
    target_regex（Optional[str]）：用于匹配目标模块的正则表达式。默认值为 None。
    modules_to_save（List[str]）：待保存的模块列表。默认值为空列表。

    lora_rank（int）：LoRA 的排名。默认值为 8。
    lora_alpha（int）：LoRA 的 Alpha 值。默认值为 32。
    lora_dropout（float）：LoRA 的丢包率。默认值为 0.05。
    允许的值为“none”和“all”。
    lora_dtype（Literal）：LoRA 的数据类型。默认值为“AUTO”。允许的值为“fp16”、“bf16”、“fp32”。 'AUTO'。
    lorap_lr_ratio（浮点数）：LoRA 的学习率。默认值为 None。
    use_rslora（布尔值）：用于指示是否使用 RSLora 的标志。默认值为 False。
    use_dora（布尔值）：用于指示是否使用 Dora 的标志。默认值为 False。
    init_weights（字符串）：支持的 Tuner 权重的初始化方法。默认值为 true。
    lora_ga_batch_size（整数）：LoRA-GA 初始化过程中用于估计梯度的批次大小。默认值为 2。
    lora_ga_iters（整数）：LoRA-GA 初始化过程中用于估计梯度的迭代次数。默认值为 2。
    lora_ga_max_length（整数）：LoRA-GA 初始化过程中用于估计梯度的最大输入长度。默认值为 1024。
    lora_ga_direction（字符串）：LoRA-GA 初始化过程中用于估计梯度的初始方向。 LoRA-GA。默认值为“ArB2r”。允许的值：‘ArBr’、‘A2rBr’、‘ArB2r’ 和 ‘random’。
    lora_ga_scale（字符串）：LoRA-GA 中初始化的缩放方法。默认值为“stable”。允许的值包括：‘gd’、‘unit’、‘stable’ 和 ‘weightS’。
    lora_ga_stable_gamma（整数）：选择“stable”缩放进行初始化时的 gamma 值。默认值为 16。

    fourier_n_frequency（整数）：FourierFT 的频率数量。默认值为 2000。
    fourier_scaling（浮点数）：FourierFT 的缩放因子。默认值为 300.0。

    boft_block_size（整数）：BOFT 的块大小。默认值为 4。
    boft_block_num（整数）：BOFT 的块数。默认值为 0。
    boft_n_butterfly_factor（整数）：BOFT 的蝴蝶因子。默认值为 1。
    boft_dropout（浮点数）：BOFT 的 Dropout 率。默认值为 0.0。

    vera_rank（整数）：Vera 的排名。默认值为 256。
    vera_projection_prng_key（整数）：Vera 投影的 PRNG 密钥。默认值为 0。
    vera_dropout（浮点数）：Vera 的 Dropout 率。默认值为 0.0。
    vera_d_initial（浮点数）：Vera D 的初始值。默认值为 0.1。

    adapter_act（字符串）：适配器的激活函数。默认值为 'gelu'。
    adapter_length（整数）：适配器的长度。默认值为 128。

    use_galore（布尔值）：用于指示是否使用 Galore 的标志。默认值为 False。
    galore_target_modules（Optional[List[str]]）：列表Galore 的目标模块。默认值为 None。
    galore_rank（整数）：Galore 的排名。默认值为 128。
    galore_update_proj_gap（整数）：更新 Galore 的投影间隙。默认值为 50。
    galore_scale（浮点数）：Galore 的缩放因子。默认值为 1.0。
    galore_proj_type（字符串）：Galore 的投影类型。默认值为 'std'。
    galore_optim_per_parameter（布尔值）：用于指示 Galore 是否针对每个参数进行优化的标志。默认值为 False。
    galore_with_embedding（布尔值）：用于指示 Galore 是否使用嵌入的标志。默认值为 False。
    galore_quantization（布尔值）：用于指示是否使用 Q-Galore 的标志。默认值为 False。
    galore_proj_quant（布尔值）：用于指示 Galore 是否使用投影量化的标志。默认值为 False。
    galore_proj_bits（整数）：投影量化的位数。默认值为 4。
    galore_proj_group_size（整数）：投影量化的组大小。默认值为 256。
    galore_cos_threshold（浮点数）：投影量化的余弦阈值。默认值为 0.4。
    galore_gamma_proj（整数）：投影量化的 Gamma。默认值为 2。
    galore_queue_size（整数）：投影量化的队列大小。默认值为 5。

    adalora_target_r（整数）：AdaLoRA 的目标等级。默认值为 8。
    adalora_init_r（整数）：AdaLoRA 的初始等级。默认值为 12。
    adalora_tinit（整数）：AdaLoRA 的初始 T 值。默认值为 100。
    adalora_tfinal（整数）：AdaLoRA 的最终 T 值。默认值为1000.
    adalora_deltaT（整数）：AdaLoRA 的 Delta T 值。默认值为 10。
    adalora_beta1（浮点数）：AdaLoRA 的 Beta1 值。默认值为 0.85。
    adalora_beta2（浮点数）：AdaLoRA 的 Beta2 值。默认值为 0.85。
    adalora_orth_reg_weight（浮点数）：正交正则化AdaLoRA 的离子权重。默认值为 0.5。

    llamapro_num_new_blocks（整数）：LLaMAPro 的新块数量。默认值为 4。

    llamapro_num_groups（可选[int]）：LLaMAPro 的组数量。默认值为 None。

    lisa_activated_layers（整数）：LISA 的激活层数量。默认值为 0。
    lisa_step_interval（整数）：LISA 激活的步长间隔。默认值为 20。

    reft_layer_key（可选[str]）：ReFT 层的密钥标识符。默认值为 None。
    reft_layers（可选[List[int]]）：参与 ReFT 的层列表。默认值为 None。
    reft_rank（整数）：ReFT 的等级参数。默认值为 4。
    reft_intervention_type（字面量）：ReFT 的干预类型。默认值为“LoreftIntervention”。
    reft_args（可选[str]）：ReFT 的附加参数。默认为“无”。
    """
    # 用于冻结部分模型参数的字段
    freeze_parameters: List[str] = field(default_factory=list)
    freeze_parameters_ratio: float = 0.  # 冻结参数的比例，范围 0~1
    trainable_parameters: List[str] = field(default_factory=list)  # 额外可训练参数

    # 是否冻结各部分模块
    freeze_llm: bool = False  # 冻结语言模型部分
    freeze_vit: bool = True  # 冻结视觉模型部分
    freeze_aligner: bool = True  # 冻结对齐模块

    # 指定要进行微调的模块，如 'all-linear'
    target_modules: List[str] = field(default_factory=lambda: ['all-linear'])
    target_regex: Optional[str] = None  # 支持使用正则表达式筛选目标模块
    modules_to_save: List[str] = field(default_factory=list)  # 指定训练时要额外保存的模块

    # LoRA 相关参数
    lora_rank: int = 8
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    lora_bias: Literal['none', 'all'] = 'none'
    lora_dtype: Literal['float16', 'bfloat16', 'float32', None] = None
    lorap_lr_ratio: Optional[float] = None
    use_rslora: bool = False  # 是否启用 RSLora
    use_dora: bool = False  # 是否启用 DoRA

    # LoRA-GA 初始化参数
    lora_ga_batch_size: int = 2
    lora_ga_iters: int = 2
    lora_ga_max_length: int = 1024
    lora_ga_direction: str = 'ArB2r'  # 初始化方向
    lora_ga_scale: str = 'stable'  # 初始化缩放方式
    lora_ga_stable_gamma: int = 16  # `stable` 模式下的 gamma 值

    # 权重初始化策略
    init_weights: str = 'true'  # 支持 'true'/'false' 字符串，也会被转换为 bool 类型

    # FourierFT 参数
    fourier_n_frequency: int = 2000
    fourier_scaling: float = 300.0

    # BOFT（块结构压缩）相关参数
    boft_block_size: int = 4
    boft_block_num: int = 0
    boft_n_butterfly_factor: int = 1
    boft_dropout: float = 0.0

    # Vera 低秩压缩相关参数
    vera_rank: int = 256
    vera_projection_prng_key: int = 0
    vera_dropout: float = 0.0
    vera_d_initial: float = 0.1

    # Adapter 相关配置
    adapter_act: str = 'gelu'
    adapter_length: int = 128

    # Galore 参数配置
    use_galore: bool = False
    galore_target_modules: Optional[List[str]] = None
    galore_rank: int = 128
    galore_update_proj_gap: int = 50
    galore_scale: float = 1.0
    galore_proj_type: str = 'std'
    galore_optim_per_parameter: bool = False
    galore_with_embedding: bool = False
    galore_quantization: bool = False
    galore_proj_quant: bool = False
    galore_proj_bits: int = 4
    galore_proj_group_size: int = 256
    galore_cos_threshold: float = 0.4
    galore_gamma_proj: int = 2
    galore_queue_size: int = 5

    # AdaLoRA 参数
    adalora_target_r: int = 8
    adalora_init_r: int = 12
    adalora_tinit: int = 0
    adalora_tfinal: int = 0
    adalora_deltaT: int = 1
    adalora_beta1: float = 0.85
    adalora_beta2: float = 0.85
    adalora_orth_reg_weight: float = 0.5

    # LLaMAPro 参数
    llamapro_num_new_blocks: int = 4
    llamapro_num_groups: Optional[int] = None

    # LISA 激活控制参数
    lisa_activated_layers: int = 0
    lisa_step_interval: int = 20

    # ReFT 插值微调方法参数
    reft_layer_key: Optional[str] = None
    reft_layers: Optional[List[int]] = None
    reft_rank: int = 4
    reft_intervention_type: Literal[
        'NoreftIntervention',
        'LoreftIntervention',
        'ConsreftIntervention',
        'LobireftIntervention',
        'DireftIntervention',
        'NodireftIntervention'
    ] = 'LoreftIntervention'
    reft_args: Optional[str] = None

    # 初始化时执行的钩子函数
    def __post_init__(self):
        # 字符串 'true'/'false' 转为布尔值
        if isinstance(self.init_weights, str) and self.init_weights.lower() in {'true', 'false'}:
            self.init_weights = bool(strtobool(self.init_weights))

        # 根据模型结构设置冻结模块等
        self._init_multimodal_full()

        # 如果设置了 target_regex，则用正则字符串替代 target_modules（可能是bug）
        if self.target_regex:
            self.target_modules = self.target_regex

    # 根据模型类型决定哪些模块冻结或训练
    def _init_multimodal_full(self):
        model_arch = get_model_arch(self.model_meta.model_arch)

        # 如果不是多模态或未能获取结构信息，则不处理
        if not self.model_meta.is_multimodal or not model_arch:
            return

        # 根据配置冻结语言模型
        if self.freeze_llm:
            self.freeze_parameters += model_arch.language_model

        # 冻结视觉模型
        if self.freeze_vit:
            self.freeze_parameters += model_arch.vision_tower

        # 冻结对齐器或加入可训练参数列表
        if self.freeze_aligner:
            self.freeze_parameters += model_arch.aligner
        else:
            self.trainable_parameters += model_arch.aligner

        # 默认冻结生成器部分
        self.freeze_parameters += model_arch.generator

        # 打印冻结与训练的参数信息
        if self.freeze_parameters:
            logger.info(f'freeze_parameters: {self.freeze_parameters}')
        if self.trainable_parameters:
            logger.info(f'additional trainable_parameters: {self.trainable_parameters}')
